# -*- coding: utf-8 -*-
"""
core/rag_retrieval/retriever.py

RAG辅助分析与代码定位模块。
负责根据需求差异，在“需求联系代码的拆解文档”和FAISS代码向量库中检索相关信息，
并整合这些信息为后续的LLM调用提供上下文。
"""

import os
import json
import faiss
import numpy as np
from core.utils import logger, read_json_file, read_file_content
# 模拟文本嵌入模型 (与vectorizer.py中的保持一致或从公共模块导入)
SIMULATED_EMBEDDING_DIM = 128 # 假设嵌入向量的维度

def get_simulated_text_embedding(text_content):
    """
    模拟获取文本内容的嵌入向量。
    与vectorizer.py中的版本保持一致。
    """
    if not text_content:
        return np.zeros(SIMULATED_EMBEDDING_DIM, dtype=np.float32)
    seed = hash(text_content) % (2**32 -1)
    np.random.seed(seed)
    vector = np.random.rand(SIMULATED_EMBEDDING_DIM).astype(np.float32)
    norm_factor = min(len(text_content) / 100.0, 1.0)
    vector = vector * norm_factor
    return vector

class RagRetriever:
    """
    RAG检索器类，封装了文档检索和代码库检索的功能。
    """
    def __init__(self, faiss_index_path, faiss_metadata_path, disassembly_doc_path, all_parsed_code_fragments_path):
        """
        初始化检索器。

        参数:
            faiss_index_path (str): FAISS索引文件路径。
            faiss_metadata_path (str): FAISS元数据文件路径。
            disassembly_doc_path (str): “需求联系代码的拆解文档”路径。
            all_parsed_code_fragments_path (str): 包含所有已解析代码片段的JSON文件路径 (用于通过索引找回完整片段)。
        """
        self.faiss_index = None
        self.faiss_metadata = None
        self.disassembly_doc_content = None
        self.all_code_fragments = None # 存储所有解析后的代码片段

        # 加载FAISS索引和元数据
        try:
            if os.path.exists(faiss_index_path):
                self.faiss_index = faiss.read_index(faiss_index_path)
                logger.info(f"FAISS索引已加载: {faiss_index_path}, 包含 {self.faiss_index.ntotal} 个向量。")
            else:
                logger.error(f"FAISS索引文件未找到: {faiss_index_path}")
            
            self.faiss_metadata = read_json_file(faiss_metadata_path)
            if self.faiss_metadata is not None:
                logger.info(f"FAISS元数据已加载: {faiss_metadata_path}, 包含 {len(self.faiss_metadata)} 条记录。")
            else:
                 logger.error(f"FAISS元数据文件加载失败或为空: {faiss_metadata_path}")
        except Exception as e:
            logger.error(f"加载FAISS索引或元数据失败: {e}")
            self.faiss_index = None # 确保失败时为None
            self.faiss_metadata = None

        # 加载“需求联系代码的拆解文档”
        self.disassembly_doc_content = read_file_content(disassembly_doc_path)
        if self.disassembly_doc_content:
            logger.info(f"需求联系代码的拆解文档已加载: {disassembly_doc_path}")
        else:
            logger.error(f"需求联系代码的拆解文档加载失败: {disassembly_doc_path}")

        # 加载所有解析后的代码片段 (用于通过FAISS元数据中的original_index找回完整信息)
        self.all_code_fragments = read_json_file(all_parsed_code_fragments_path)
        if self.all_code_fragments:
            logger.info(f"所有已解析的代码片段已加载: {all_parsed_code_fragments_path}, 共 {len(self.all_code_fragments)} 个片段。")
        else:
            logger.error(f"所有已解析的代码片段加载失败: {all_parsed_code_fragments_path}")

    def retrieve_from_disassembly_doc(self, query_text, top_k=2):
        """
        从“需求联系代码的拆解文档”中检索相关片段。
        简单实现：基于关键词匹配或返回固定内容（模拟）。
        实际应用中，这里可以使用更高级的文本检索方法，如TF-IDF、BM25或语义检索。

        参数:
            query_text (str): 查询文本 (例如，需求差异的描述)。
            top_k (int): 返回最相关的片段数量。

        返回:
            list: 包含相关文本片段的列表，每个片段是一个字典 {"title": ..., "content": ..., "distance/score": ...}。
        """
        logger.info(f"开始从拆解文档中检索，查询: 	'{query_text[:50]}...' ")
        if not self.disassembly_doc_content:
            logger.warning("拆解文档内容为空，无法进行检索。")
            return []

        # --- 模拟检索逻辑 ---
        # 这是一个非常基础的模拟，实际中需要替换为真正的检索算法
        results = []
        # 假设文档是以某种方式结构化的，例如有标题和段落
        # 或者像用户提供的查询结果那样，有标题、距离、内容
        # 这里我们简单地返回文档的前几行或包含查询关键词的行
        
        # 模拟用户提供的查询结果格式
        # 实际中，这部分应该由一个独立的文本检索引擎（如Whoosh, Elasticsearch, 或基于向量的）完成
        if "童锁" in query_text:
            results.append({
                "title": "拆解文档中关于童锁的说明",
                "distance": 0.1,
                "content": "童锁功能通常涉及按键模块 (key.c) 用于检测特定组合键长按，以及LED模块 (LED.c) 用于状态指示。可能还需要在任务调度 (task.c) 中处理状态机。"
            })
        if "电量显示" in query_text:
            results.append({
                "title": "拆解文档中关于电量显示的说明",
                "distance": 0.15,
                "content": "电量显示逻辑一般在ADC采样后，由特定函数处理电压到显示的映射，可能在 modular.c 或专门的 power_management.c 中。LED的控制在LED.c。"
            })
        
        # 如果没有特定匹配，可以返回一个通用提示或文档摘要
        if not results:
            results.append({
                "title": "拆解文档通用参考",
                "distance": 0.5,
                "content": self.disassembly_doc_content[:500] + "... (文档摘要)"
            })
        
        logger.info(f"从拆解文档检索到 {len(results)} 个结果。")
        return results[:top_k]

    def retrieve_from_codebase(self, query_text, top_k=3):
        """
        从FAISS代码向量库中检索相关的代码片段。

        参数:
            query_text (str): 查询文本 (例如，需求差异的描述)。
            top_k (int): 返回最相关的代码片段数量。

        返回:
            list: 包含相关代码片段信息的列表。每个元素是原始代码片段字典。
        """
        logger.info(f"开始从代码库FAISS索引中检索，查询: 	'{query_text[:50]}...' ")
        if not self.faiss_index or self.faiss_metadata is None or self.all_code_fragments is None:
            logger.warning("FAISS索引、元数据或完整代码片段列表未加载，无法进行代码库检索。")
            return []

        query_vector = get_simulated_text_embedding(query_text)
        query_vector_np = np.array([query_vector]).astype("float32")

        try:
            distances, indices = self.faiss_index.search(query_vector_np, top_k)
            retrieved_fragments = []
            for i in range(len(indices[0])):
                idx = indices[0][i]
                dist = distances[0][i]
                if idx < 0 or idx >= len(self.faiss_metadata): # 索引可能无效
                    logger.warning(f"检索到无效的FAISS索引: {idx}")
                    continue
                
                metadata_item = self.faiss_metadata[idx]
                original_fragment_index = metadata_item.get("original_index")
                
                if original_fragment_index is not None and 0 <= original_fragment_index < len(self.all_code_fragments):
                    full_fragment = self.all_code_fragments[original_fragment_index]
                    # 可以将距离也附加到结果中，如果需要
                    full_fragment["retrieval_distance"] = float(dist) 
                    retrieved_fragments.append(full_fragment)
                else:
                    logger.warning(f"无法从元数据中找到有效的原始片段索引: {original_fragment_index} (FAISS idx: {idx})")
            
            logger.info(f"从代码库FAISS索引检索到 {len(retrieved_fragments)} 个代码片段。")
            return retrieved_fragments
        except Exception as e:
            logger.error(f"FAISS检索失败: {e}")
            return []

    def retrieve_context(self, requirement_diff_item):
        """
        针对单个需求差异项，进行多路检索并整合上下文。

        参数:
            requirement_diff_item (dict): 单个需求差异项的字典，
                                        例如: {"type": "新增", "feature_name": "童锁功能", ...}

        返回:
            dict: 包含检索到的文档片段和代码片段的字典。
                  例如: {"related_docs": [...], "related_code": [...]}
        """
        query = requirement_diff_item.get("feature_name", "") + " " + \ 
                requirement_diff_item.get("new_description", "") + " " + \ 
                requirement_diff_item.get("details", "")
        query = query.strip()
        if not query:
            logger.warning("无法从需求差异项构建查询文本。")
            return {"related_docs": [], "related_code": []}

        logger.info(f"开始为需求 '{requirement_diff_item.get("feature_name")}' 整合RAG上下文...")
        
        related_docs = self.retrieve_from_disassembly_doc(query)
        related_code = self.retrieve_from_codebase(query)

        return {
            "related_docs": related_docs,
            "related_code": related_code
        }

if __name__ == "__main__":
    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))

    # --- 模拟输入文件路径 (这些文件应由之前的模块生成) ---
    mock_faiss_index_file = os.path.join(project_root, "data", "output", "faiss_index", "codebase.index")
    mock_faiss_metadata_file = os.path.join(project_root, "data", "output", "faiss_index", "metadata.json")
    mock_disassembly_doc = os.path.join(project_root, "data", "input", "disassembly_doc.md")
    mock_all_fragments_file = os.path.join(project_root, "data", "output", "parsed_code_fragments.json")

    # --- 确保模拟文件存在 (通常由其他脚本创建，这里为了独立测试) ---
    # 假设 vectorizer.py 和 preprocessor.py 已经运行并生成了这些文件
    # 如果没有，这个测试会失败或返回空结果
    if not os.path.exists(mock_faiss_index_file) or \ 
       not os.path.exists(mock_faiss_metadata_file) or \ 
       not os.path.exists(mock_all_fragments_file):
        logger.warning("一个或多个必需的输入文件 (FAISS索引/元数据/解析片段) 不存在。RAG检索器测试可能不完整。")
        logger.warning(f"请确保运行 preprocessor.py 和 vectorizer.py 以生成这些文件:")
        logger.warning(f"  - {mock_all_fragments_file}")
        logger.warning(f"  - {mock_faiss_index_file}")
        logger.warning(f"  - {mock_faiss_metadata_file}")

    # 创建一个模拟的拆解文档
    os.makedirs(os.path.dirname(mock_disassembly_doc), exist_ok=True)
    with open(mock_disassembly_doc, "w", encoding="utf-8") as f_dis:
        f_dis.write("""
        # 需求联系代码的拆解文档 (模拟)

        ## 童锁功能实现思路
        童锁功能通常涉及按键模块 (key.c) 用于检测特定组合键长按，以及LED模块 (LED.c) 用于状态指示。
        可能还需要在任务调度 (task.c) 中处理状态机，以管理童锁激活状态。
        相关函数可能包括 `Key_Get_LongPress_Event()` 和 `LED_Set_ChildLock_Indicator()`。

        ## 电量显示逻辑
        电量显示逻辑一般在ADC采样后，由特定函数处理电压到显示的映射。
        此映射逻辑可能在 `modular.c` 或专门的 `power_management.c` 中。
        LED的控制在 `LED.c` 中的 `LED_Display_Battery()` 函数。
        注意ADC初始化的代码在 `init.c`。
        """)

    # --- 初始化RAG检索器 ---
    retriever = RagRetriever(mock_faiss_index_file, 
                             mock_faiss_metadata_file, 
                             mock_disassembly_doc, 
                             mock_all_fragments_file)

    # --- 模拟一个需求差异项进行测试 ---
    mock_requirement_item_1 = {
        "type": "新增",
        "feature_name": "童锁保护功能",
        "new_description": "在新需求中，增加了童锁保护功能。用户可以通过长按某个预设的组合键3秒来激活或解除童锁模式。",
        "old_description": "",
        "details": "此功能为全新加入，旨在提升产品在有儿童环境下的使用安全性。"
    }
    
    logger.info("\n--- 测试RAG检索: 童锁保护功能 ---")
    context1 = retriever.retrieve_context(mock_requirement_item_1)
    logger.info("相关文档片段:")
    for doc in context1["related_docs"]:
        logger.info(f"  - 标题: {doc.get('title')}, 内容: {doc.get('content', '')[:100]}...")
    logger.info("相关代码片段:")
    for code in context1["related_code"]:
        logger.info(f"  - 文件: {code.get('file_path')}, 名称: {code.get('name')}, 类型: {code.get('type')}, 距离: {code.get('retrieval_distance'):.4f}")
        logger.info(f"    代码: {code.get('code', '')[:100].replace('\n', ' ')}...")

    mock_requirement_item_2 = {
        "type": "修改",
        "feature_name": "电量显示逻辑",
        "new_description": "新需求调整了电量显示方式。具体为：当电池电压低于3.2V时，LED显示1格电量。",
        "old_description": "旧版需求中，电量显示可能为多级LED指示。",
        "details": "修改的目的是为了简化电量指示。"
    }
    logger.info("\n--- 测试RAG检索: 电量显示逻辑 ---")
    context2 = retriever.retrieve_context(mock_requirement_item_2)
    logger.info("相关文档片段:")
    for doc in context2["related_docs"]:
        logger.info(f"  - 标题: {doc.get('title')}, 内容: {doc.get('content', '')[:100]}...")
    logger.info("相关代码片段:")
    for code in context2["related_code"]:
        logger.info(f"  - 文件: {code.get('file_path')}, 名称: {code.get('name')}, 类型: {code.get('type')}, 距离: {code.get('retrieval_distance'):.4f}")
        logger.info(f"    代码: {code.get('code', '')[:100].replace('\n', ' ')}...")

    logger.info("\nRAG检索器模块独立测试完成。")


